import os
import click
import requests
import multiprocessing
from rich.console import Console
from rich.table import Table
from .config import Config

console = Console()
config = Config.load()

STATUS_COLORS = {
    "PENDING": "yellow",
    "RUNNING": "green",
    "COMPLETED": "blue",
    "FAILED": "red"
}
# 后端服务的URL
IP = "127.0.0.1"
PORT = 12030
BASE_URL = f"http://{IP}:{PORT}"  # 根据实际情况调整

def server_health() -> bool:
    try:
        # 如果'/health'返回200，则服务正常
        response = requests.get(BASE_URL + "/health", timeout=0.5)
        if response.status_code == 200:
            return True
    except requests.exceptions.RequestException as e:
        #console.print(f"[red]Error checking server health: {e}[/red]")
        return False
    return False
        
def print_audit_log(ctx, param, value):
    if not value or ctx.resilient_parsing:
        return
    console.print(f"[yellow]审计日志: {' '.join(ctx.args)}[/yellow]")
    return value

def get_server_execution_file_path():
    import sys
    if getattr(sys, 'frozen', False): # 程序被打包运行
        return [sys.executable, "--server"]
    else:
        # 如果不是打包运行，则使用当前目录
        base_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        server_file = os.path.join(base_path, "scheduler.py")
        return [sys.executable, server_file, "--server"]

def start_server(config, ip: str = "127.0.0.1", port: int = 12030, **kwargs):
    import subprocess
    from datetime import datetime
    # 创建log文件
    os.makedirs(config.log_dir, exist_ok=True)
    log_file = os.path.join(config.log_dir, f"scheduler-server-{datetime.now().strftime('%Y%m%d-%H%M%S')}.log")
    # 复制环境变量
    console.print(f"[yellow]启动后端服务: {' '.join(get_server_execution_file_path())}[/yellow]")
    env = os.environ.copy()
    with open(log_file, "w") as f:
        p = subprocess.Popen(get_server_execution_file_path() + ["--host", ip, "--port", str(port)],
                         stdout=f, 
                         stderr=f,
                         cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                         env=env,
                         )
        config.server_pid = p.pid
        config.save()
        
        
@click.group()
@click.option("--verbose", is_flag=True, callback=print_audit_log, expose_value=False, default=False)
def cli():
    """深度学习任务调度器命令行接口"""

@cli.command()
def start():
    """启动后端服务"""
    try:
        # 如果后端服务开启，则退出
        if server_health():
            console.print("[red]× 后端服务已开启，无需重复开启服务[/red]")
            return
        
        # 创建并启动新进程
        if os.name == 'nt':
            multiprocessing.freeze_support()
        p = multiprocessing.Process(target=start_server,
                                    name="task-scheduler-server",  # 设置进程名称
                                    daemon=False,  # 确保即使主进程退出，子进程也继续运行
                                    args=(config, IP, PORT)
                                    ) 
        p.start()

        console.print(f"[green]✓ 后端服务已启动, pid:{p.pid}[/green]")
    except Exception as e:
        console.print(f"[red]× 启动后端服务失败: {str(e)}[/red]")
        
        
@cli.command()
@click.argument("command")
@click.option("-d", "--device", type=str, help="GPU设备名,如cuda:0,cuda:1,多个设备用逗号分隔")
@click.option("-g", "--num-gpu", default=0, type=int, help="申请GPU个数")  # 默认为0，表示不需要GPU
@click.option("-c", "--mem-cpu", type=float, help="申请CPU内存") 
@click.option("-m", "--mem-per-gpu", type=float, help="单个GPU需要申请的内存")
@click.option("--max-retries", type=int)
def add_task(command, device, num_gpu, mem_cpu, mem_per_gpu, max_retries):
    """添加一个任务"""
    url = f"{BASE_URL}/add-task"
    payload = {
        "command": command,
        "device": list(device.split(",")),
        "num_gpu": num_gpu,
        "mem_cpu": mem_cpu,
        "mem_per_gpu": mem_per_gpu,
        "env": os.environ.copy(),
        "pwd": os.getcwd(),
    }
    if max_retries is not None:
        payload["max_retries"] = max_retries
    try:
        response = requests.post(url, json=payload)
        if response.status_code == 200:
            console.print(f"[green]✓ 任务已添加 (ID: {response.json().get('message').split(': ')[-1]})[/green]")
        else:
            console.print(f"[red]× 添加任务失败: {response.json().get('error')}[/red]")
    except Exception as e:
        console.print(f"[red]× 发生错误: {str(e)}[/red]")

@cli.command()
def status():
    """查看任务和设备状态"""
    url = f"{BASE_URL}/status"
    try:
        response = requests.get(url)
        if response.status_code == 200:
            data = response.json()
            
            gpu_table = Table(title="GPU 资源状态")
            gpu_table.add_column("Device", style="cyan")
            gpu_table.add_column("Used/Total", justify="right")
            for gpu in data["gpu_status"]:
                gpu_table.add_row(
                    gpu["device"],
                    f"{gpu['used']} / {gpu['total']} GB",
                    style="red" if float(gpu['used']) > float(gpu['total']) * 0.9 else "green"
                )
            
            cpu_data = data["cpu_status"]
            cpu_table = Table(title="CPU 资源状态")
            cpu_table.add_column("Total Memory", style="cyan")
            cpu_table.add_column("Reserved Memory", style="cyan")
            cpu_table.add_column("Available Memory", style="cyan")
            cpu_table.add_column("Used Memory", justify="right")
            cpu_table.add_row(
                str(cpu_data["total_mem"]),
                str(cpu_data["reserved_mem"]),
                str(cpu_data["available_mem"]),
                str(cpu_data["used_mem"]),
                style="red" if cpu_data["used_mem"] > cpu_data["available_mem"] * 0.9 else "green"
            )
            
            task_table = Table(title="任务队列")
            task_table.add_column("ID", style="magenta")
            task_table.add_column("Status", width=12)
            task_table.add_column("Command", width=40)
            task_table.add_column("GPUs")
            task_table.add_column("CPU Memory")
            task_table.add_column("Retries")
            for task in data["tasks"]:
                task_table.add_row(
                    task["id"],
                    f"[{STATUS_COLORS[task['status']]}]{task['status']}[/]",
                    task["command"][:40] + ("..." if len(task["command"]) > 40 else ""),
                    task["gpus"] if task["gpus"] else "-",
                    str(task["cpu-mem"]),
                    f"{task['retries'].split('/')[0]}/{task['retries'].split('/')[1]}"
                )
            
            console.print(gpu_table)
            console.print(cpu_table)
            console.print(task_table)
        else:
            console.print(f"[red]× 获取状态失败: {response.json().get('error')}[/red]")
    except Exception as e:
        console.print(f"[red]× 发生错误: {str(e)}[/red]")

@cli.command()
@click.argument("task_id")
def logs(task_id):
    """查看任务日志"""
    url = f"{BASE_URL}/logs/{task_id}"
    try:
        response = requests.get(url, timeout=10)  # 设置全局请求超时
        if response.status_code == 200:
            data = response.json()
            console.print(f"[bold]任务日志 {data['file_name']}:[/bold]")
            console.print(data['log_content'])
        elif response.status_code == 423:  # 明确处理锁定状态
            console.print(f"[yellow]! 日志文件被其他进程锁定，请稍后重试[/yellow]")
        elif response.status_code == 404:
            console.print(f"[yellow]! {response.json().get('error')}[/yellow]")
        elif response.status_code == 400:
            data = response.json()
            console.print(f"[yellow]! {data['error']}[/yellow]")
            console.print("匹配的任务ID列表:")
            for log in data.get('matched_task_logs', []):
                console.print(f"  - {log}")
        else:
            console.print(f"[red]× 错误: HTTP {response.status_code} - {response.text}[/red]")
    except requests.exceptions.Timeout:
        console.print("[red]× 请求超时，请检查网络连接[/red]")
    except Exception as e:
        console.print(f"[red]× 未知错误: {str(e)}[/red]")


@cli.command()
@click.option("--available-devices", help="可用设备列表，逗号分隔")
@click.option("--reserved-mem", type=int, help="每GPU预留内存(MB)")
@click.option("--reserved-cpu-mem", type=int, help="预留CPU内存(MB)")
@click.option("--max-retries", type=int, help="最大重试次数")
def config_update(available_devices, reserved_mem, reserved_cpu_mem, max_retries):
    """更新调度器配置"""
    url = f"{BASE_URL}/config-update"
    payload = {}
    if available_devices:
        payload["available_devices"] = available_devices
    if reserved_mem is not None:
        payload["reserved_mem"] = reserved_mem
    if reserved_cpu_mem is not None:
        payload["reserved_cpu_mem"] = reserved_cpu_mem
    if max_retries is not None:
        payload["max_retries"] = max_retries
    
    try:
        response = requests.post(url, json=payload)
        if response.status_code == 200:
            console.print("[green]✓ 配置已更新[/green]")
            console.print(response.json().get("config"))
        else:
            console.print(f"[red]× 配置更新失败: {response.json().get('error')}[/red]")
    except Exception as e:
        console.print(f"[red]× 发生错误: {str(e)}[/red]")

@cli.command()
def stop():
    """停止调度器"""
    url = f"{BASE_URL}/stop"
    try:
        response = requests.post(url)
        if response.status_code == 200:
            console.print("[green]✔ 停止调度器成功[/green]")
        else:
            console.print(f"[red]× 停止调度器失败: {response.json().get('error')}[/red]")
    except Exception as e:
        console.print(f"[red]× 发生错误: {str(e)}[/red]")

if __name__ == "__main__":
    cli()
